package com.btl.service.ai.support.rag.scoring.strategy;

import com.btl.service.ai.support.rag.embedding.strategy.RetrieverStrategy;
import com.btl.service.ai.support.rag.scoring.ScoringModelEnum;
import com.btl.service.ai.support.rag.scoring.ScoringStrategy;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import dev.langchain4j.rag.query.Query;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;

@Service
public class BegScoringBaseScoringModelStrategy implements ScoringStrategy {

    private final ScoringModel begRerankBaseRerankModel;

    public BegScoringBaseScoringModelStrategy(@Qualifier("begRerankBaseScoringModel") ScoringModel begRerankBaseRerankModel) {
        this.begRerankBaseRerankModel = begRerankBaseRerankModel;
    }

    @Override
    public List<Content> rerank(List<RetrieverStrategy> retrieverStrategies, String collectionName, Query query) throws InterruptedException {
        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
            List<Callable<List<Content>>> tasks = retrieverStrategies.stream()
                    .map(r -> (Callable<List<Content>>) () -> r.retriever(collectionName, query))
                    .toList();
            List<Content> allSegments = executor.invokeAll(tasks)
                    .stream()
                    .flatMap(f -> {
                        try {
                            return f.get()
                                    .stream();
                        } catch (Exception e) {
                            throw new RuntimeException(e);
                        }
                    })
                    .toList();
            Function<Map<Query, Collection<List<Content>>>, Query> querySelector = queryMap -> query;
            ReRankingContentAggregator contentAggregator = new ReRankingContentAggregator(begRerankBaseRerankModel, querySelector, 0.75);
            // 进行聚合
            Map<Query, Collection<List<Content>>> queryToContents = new HashMap<>();
            queryToContents.put(query, List.of(allSegments));
            List<Content> ranked = contentAggregator.aggregate(queryToContents);
            // 打印结果
            for (int i = 0; i < ranked.size(); i++) {
                System.out.println("Rank " + (i + 1) + ": " + ranked.get(i)
                        .textSegment());
            }
            return ranked;
        }
    }

    @Override
    public ScoringModelEnum model() {
        return ScoringModelEnum.LOCAL_BGE_RERANK_SCORING_MODEL;
    }
}
